// Copyright (c) 2010, Suryandaru Triandana. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This package implements the Salsa20 encryption algorithm and its variant Salsa20/8 and Salsa20/12 as defined in eSTREAM portfolio.
package salsa20

// This Go implementation is derived in part from the reference
// ANSI C implementation, which carries the following notice:
//
//	salsa20-ref.c version 20051118
//	D. J. Bernstein
//	Public domain.

import (
	"os"
	"strconv"
)

// A Cipher is an instance of Salsa20 encryption using a particular key.
type Cipher struct {
	i [16]uint32
	r []byte
	round byte
}

type KeySizeError int

func (k KeySizeError) String() string {
	return "crypto/salsa20: invalid key size " + strconv.Itoa(int(k))
}

type IVSizeError int

func (k IVSizeError) String() string {
	return "crypto/salsa20: invalid iv size " + strconv.Itoa(int(k))
}

var (
	sigma = [4]uint32{0x61707865, 0x3320646e, 0x79622d32, 0x6b206574}
	tau = [4]uint32{0x61707865, 0x3120646e, 0x79622d36, 0x6b206574}
)

func rotl(v, n uint32) uint32 {
	return v<<n | v>>(32-n)
}

func newCipher(key []byte, round byte) (*Cipher, os.Error) {
	k := len(key)
	if k != 16 && k != 32 {
		return nil, KeySizeError(k)
	}
	var c Cipher
	c.round = round
	c.i[1] = uint32(key[ 0]) | uint32(key[ 1])<<8 | uint32(key[ 2])<<16 | uint32(key[ 3])<<24
	c.i[2] = uint32(key[ 4]) | uint32(key[ 5])<<8 | uint32(key[ 6])<<16 | uint32(key[ 7])<<24
	c.i[3] = uint32(key[ 8]) | uint32(key[ 9])<<8 | uint32(key[10])<<16 | uint32(key[11])<<24
	c.i[4] = uint32(key[12]) | uint32(key[13])<<8 | uint32(key[14])<<16 | uint32(key[15])<<24
	ck := tau
	switch(k) {
	case 16:
		c.i[11] = c.i[1]
		c.i[12] = c.i[2]
		c.i[13] = c.i[3]
		c.i[14] = c.i[4]
		break
	case 32:
		ck = sigma
		c.i[11] = uint32(key[16]) | uint32(key[17])<<8 | uint32(key[18])<<16 | uint32(key[19])<<24
		c.i[12] = uint32(key[20]) | uint32(key[21])<<8 | uint32(key[22])<<16 | uint32(key[23])<<24
		c.i[13] = uint32(key[24]) | uint32(key[25])<<8 | uint32(key[26])<<16 | uint32(key[27])<<24
		c.i[14] = uint32(key[28]) | uint32(key[29])<<8 | uint32(key[30])<<16 | uint32(key[31])<<24
		break
	}
	c.i[0] = ck[0]
	c.i[5] = ck[1]
	c.i[10] = ck[2]
	c.i[15] = ck[3]
	return &c, nil
}

// NewCipher creates and returns 20-round Salsa20 Cipher.
// Salsa20 key, must be 16 or 32 bytes.
func NewCipher(key []byte) (*Cipher, os.Error) {
	return newCipher(key, 20)
}

// NewCipher12 creates and returns 8-round Salsa20/8 Cipher.
// Salsa20/8 key, must be 16 or 32 bytes.
func NewCipher8(key []byte) (*Cipher, os.Error) {
	return newCipher(key, 8)
}

// NewCipher12 creates and returns 12-round Salsa20/12 Cipher.
// Salsa20/12 key, must be 16 or 32 bytes.
func NewCipher12(key []byte) (*Cipher, os.Error) {
	return newCipher(key, 12)
}

// SetupIV will setup Initialization vector.
// Rabbit iv, must be 8 bytes.
func (c *Cipher) SetupIV(iv []byte) os.Error {
	k := len(iv)
	if k != 8 {
		return IVSizeError(k)
	}
	c.i[6] = uint32(iv[0]) | uint32(iv[1])<<8 | uint32(iv[2])<<16 | uint32(iv[3])<<24
	c.i[7] = uint32(iv[4]) | uint32(iv[5])<<8 | uint32(iv[6])<<16 | uint32(iv[7])<<24
	c.i[8] = 0
	c.i[9] = 0
	c.r = nil
	return nil
}

// ProcessStream will encrypt or decrypt given buffer.
func (c *Cipher) ProcessStream(buf []byte) {
	l := len(buf)
	if l == 0 {
		return;
	}
	i := 0
	if m := len(c.r); m > 0 {
		for ; i < m && i < l; i++ {
			buf[i] ^= c.r[i]
		}
		c.r = nil
	}
	for i < l {
		o0 , o1 , o2 , o3  := c.i[ 0], c.i[ 1], c.i[ 2], c.i[ 3]
		o4 , o5 , o6 , o7  := c.i[ 4], c.i[ 5], c.i[ 6], c.i[ 7]
		o8 , o9 , o10, o11 := c.i[ 8], c.i[ 9], c.i[10], c.i[11]
		o12, o13, o14, o15 := c.i[12], c.i[13], c.i[14], c.i[15]

		for r := c.round; r > 0; r-=2 {
			o4  = o4  ^ rotl(o0  + o12, 7)
			o8  = o8  ^ rotl(o4  + o0,  9)
			o12 = o12 ^ rotl(o8  + o4, 13)
			o0  = o0  ^ rotl(o12 + o8, 18)
			o9  = o9  ^ rotl(o5  + o1,  7)
			o13 = o13 ^ rotl(o9  + o5,  9)
			o1  = o1  ^ rotl(o13 + o9, 13)
			o5  = o5  ^ rotl(o1  + o13,18)
			o14 = o14 ^ rotl(o10 + o6,  7)
			o2  = o2  ^ rotl(o14 + o10, 9)
			o6  = o6  ^ rotl(o2  + o14,13)
			o10 = o10 ^ rotl(o6  + o2, 18)
			o3  = o3  ^ rotl(o15 + o11, 7)
			o7  = o7  ^ rotl(o3  + o15, 9)
			o11 = o11 ^ rotl(o7  + o3, 13)
			o15 = o15 ^ rotl(o11 + o7, 18)
			o1  = o1  ^ rotl(o0  + o3,  7)
			o2  = o2  ^ rotl(o1  + o0,  9)
			o3  = o3  ^ rotl(o2  + o1, 13)
			o0  = o0  ^ rotl(o3  + o2, 18)
			o6  = o6  ^ rotl(o5  + o4,  7)
			o7  = o7  ^ rotl(o6  + o5,  9)
			o4  = o4  ^ rotl(o7  + o6, 13)
			o5  = o5  ^ rotl(o4  + o7, 18)
			o11 = o11 ^ rotl(o10 + o9,  7)
			o8  = o8  ^ rotl(o11 + o10, 9)
			o9  = o9  ^ rotl(o8  + o11,13)
			o10 = o10 ^ rotl(o9  + o8, 18)
			o12 = o12 ^ rotl(o15 + o14, 7)
			o13 = o13 ^ rotl(o12 + o15, 9)
			o14 = o14 ^ rotl(o13 + o12,13)
			o15 = o15 ^ rotl(o14 + o13,18)
		}

		o0  += c.i[0]
		o1  += c.i[1]
		o2  += c.i[2]
		o3  += c.i[3]
		o4  += c.i[4]
		o5  += c.i[5]
		o6  += c.i[6]
		o7  += c.i[7]
		o8  += c.i[8]
		o9  += c.i[9]
		o10 += c.i[10]
		o11 += c.i[11]
		o12 += c.i[12]
		o13 += c.i[13]
		o14 += c.i[14]
		o15 += c.i[15]

		c.i[8] += 1
		if c.i[8] == 0 {
			c.i[9] += 1
		}

		n := l - i

		if n >= 64 {
			buf[i + 0] ^= byte(o0     )
			buf[i + 1] ^= byte(o0 >> 8)
			buf[i + 2] ^= byte(o0 >>16)
			buf[i + 3] ^= byte(o0 >>24)
			buf[i + 4] ^= byte(o1     )
			buf[i + 5] ^= byte(o1 >> 8)
			buf[i + 6] ^= byte(o1 >>16)
			buf[i + 7] ^= byte(o1 >>24)
			buf[i + 8] ^= byte(o2     )
			buf[i + 9] ^= byte(o2 >> 8)
			buf[i +10] ^= byte(o2 >>16)
			buf[i +11] ^= byte(o2 >>24)
			buf[i +12] ^= byte(o3     )
			buf[i +13] ^= byte(o3 >> 8)
			buf[i +14] ^= byte(o3 >>16)
			buf[i +15] ^= byte(o3 >>24)
			buf[i +16] ^= byte(o4     )
			buf[i +17] ^= byte(o4 >> 8)
			buf[i +18] ^= byte(o4 >>16)
			buf[i +19] ^= byte(o4 >>24)
			buf[i +20] ^= byte(o5     )
			buf[i +21] ^= byte(o5 >> 8)
			buf[i +22] ^= byte(o5 >>16)
			buf[i +23] ^= byte(o5 >>24)
			buf[i +24] ^= byte(o6     )
			buf[i +25] ^= byte(o6 >> 8)
			buf[i +26] ^= byte(o6 >>16)
			buf[i +27] ^= byte(o6 >>24)
			buf[i +28] ^= byte(o7     )
			buf[i +29] ^= byte(o7 >> 8)
			buf[i +30] ^= byte(o7 >>16)
			buf[i +31] ^= byte(o7 >>24)
			buf[i +32] ^= byte(o8     )
			buf[i +33] ^= byte(o8 >> 8)
			buf[i +34] ^= byte(o8 >>16)
			buf[i +35] ^= byte(o8 >>24)
			buf[i +36] ^= byte(o9     )
			buf[i +37] ^= byte(o9 >> 8)
			buf[i +38] ^= byte(o9 >>16)
			buf[i +39] ^= byte(o9 >>24)
			buf[i +40] ^= byte(o10    )
			buf[i +41] ^= byte(o10>> 8)
			buf[i +42] ^= byte(o10>>16)
			buf[i +43] ^= byte(o10>>24)
			buf[i +44] ^= byte(o11    )
			buf[i +45] ^= byte(o11>> 8)
			buf[i +46] ^= byte(o11>>16)
			buf[i +47] ^= byte(o11>>24)
			buf[i +48] ^= byte(o12    )
			buf[i +49] ^= byte(o12>> 8)
			buf[i +50] ^= byte(o12>>16)
			buf[i +51] ^= byte(o12>>24)
			buf[i +52] ^= byte(o13    )
			buf[i +53] ^= byte(o13>> 8)
			buf[i +54] ^= byte(o13>>16)
			buf[i +55] ^= byte(o13>>24)
			buf[i +56] ^= byte(o14    )
			buf[i +57] ^= byte(o14>> 8)
			buf[i +58] ^= byte(o14>>16)
			buf[i +59] ^= byte(o14>>24)
			buf[i +60] ^= byte(o15    )
			buf[i +61] ^= byte(o15>> 8)
			buf[i +62] ^= byte(o15>>16)
			buf[i +63] ^= byte(o15>>24)
			i += 64
			continue
		}

		j := 0
		b := buf

fillr:
		for z, y := o0, o1; j < 16 && n >= 8; n-=8 {
			switch(j) {
			case  2: z, y =  o2, o3
			case  4: z, y =  o4, o5
			case  6: z, y =  o6, o7
			case  8: z, y =  o8, o9
			case 10: z, y = o10, o11
			case 12: z, y = o12, o13
			case 14: z, y = o14, o15
			}
			b[i+0] ^= byte(z    )
			b[i+1] ^= byte(z>> 8)
			b[i+2] ^= byte(z>>16)
			b[i+3] ^= byte(z>>24)
			b[i+4] ^= byte(y    )
			b[i+5] ^= byte(y>> 8)
			b[i+6] ^= byte(y>>16)
			b[i+7] ^= byte(y>>24)
			i, j = i + 8, j + 2
		}

		if n <= 0 {
			if j >= 16 { return }
			n = (16 - j)*4
			c.r = make([]byte, n)
			b = c.r
			i = 0
			goto fillr
		}

		for z, f := o0, false; j < 16; j++ {
			switch(j) {
			case  1: z =  o1
			case  2: z =  o2
			case  3: z =  o3
			case  4: z =  o4
			case  5: z =  o5
			case  6: z =  o6
			case  7: z =  o7
			case  8: z =  o8
			case  9: z =  o9
			case 10: z = o10
			case 11: z = o11
			case 12: z = o12
			case 13: z = o13
			case 14: z = o14
			case 15: z = o15
			}
			for k := uint32(0); k < 4; k++ {
				b[i] ^= byte(z>>(k*8))
				if i++; f == false && i >= l {
					l = (15 - j)*4 + (3 - int(k))
					if l == 0 { return }
					c.r = make([]byte, l)
					b, i, f = c.r, 0, true
				}
			}
		}
	}
}

// ResetCipher resets cipher round to original state. Initialization vector will be erased.
func (c *Cipher) ResetCipher() {
	c.i[6] = 0
	c.i[7] = 0
	c.i[8] = 0
	c.i[9] = 0
	c.r = nil
}

// Reset zeros the key data so that it will no longer appear in the
// process's memory.
func (c *Cipher) Reset() {
	for i := range c.i {
		c.i[i] = 0
	}
	c.round = 0
}

